package dataframe

import org.apache.spark.SparkConf
import org.apache.spark.sql.{DataFrame, SparkSession}
import org.apache.spark.sql.functions.{col, column, expr}

object DataFrame_Operation {
  def main(args: Array[String]): Unit = {
    val conf = new SparkConf()
    conf.setMaster("local[*]")
    conf.setAppName("DataFrame_Operation")

    val spark = SparkSession
      .builder()
      .config(conf)
      .getOrCreate()
    import spark.implicits._

    val kvDF = Seq(
      (1, 2), (3, 4)
    ).toDF("key", "value")

    // 取得DataFrame中所有列
    println(kvDF.columns.mkString(","))

    // 以不同的方式选择特定的列
    kvDF.select("key").show()
    kvDF.select(col("key")).show()
    kvDF.select(column("key")).show()
    kvDF.select(expr("key")).show()
    kvDF.select($"key").show()
    kvDF.select('key).show()

    // 使用DataFrame的col函数
    kvDF.select(kvDF.col("key")).show()

    // DataFrame各种操作
    val moviesDF: DataFrame = spark
      .read
      .option("header", true)
      .option("sep", ",")
      .option("inferSchema", true)
      .csv("data/movies_DF.csv")
    moviesDF.printSchema()
    moviesDF.show()

    // select函数，选择指定的列
    moviesDF.select("title","year").show(5)
    // 使用column表达式将电影上演的年份转换到年代
    moviesDF.select(
      col("title"),
      (column("year")-column("year")%10).alias("decade")
    ).show(5)

    // selectExpr：使用SQL表达式选择列
    moviesDF.selectExpr("title","year-year%10 as decade").show(5)
    // 使用了SQL表达式和内置函数，用来查询电影数量和演员数量这两个值
    moviesDF.selectExpr(
      "count(distinct(title)) as movies_count",
      "count(distinct(actor)) as actors_count"
    ).show()

    // filter, where：实现过滤
    moviesDF.filter($"year"<2000).show(5)
    // moviesDF.where("year<2000").show()
    // 找出2000年及以后上映的电影
    moviesDF.filter($"year">=2000).show(5)
    // 找出2000年上映的电影。注意，在Scala中，相等比较要求3个等号
    moviesDF.filter($"year"===2000).show(5)
    // 找出非2000年上映的电影。注意，在Scala中，不相等比较使用的操作符是 =!=
    moviesDF.filter($"year"=!=2000).show(5)


    // 对DataFrame进行action操作
    println("*"*100)
    // 查看前5条数据。第2个参数指定当列内容较长时，是否截断显示，false为不截断
    moviesDF.show(5,false)
    // 返回数据集中的数量
    println(moviesDF.count())
    // 返回数据集中第1条数据
    println(moviesDF.first())
    println(f"第一条数据：\n演员：${moviesDF.first()(0)}，电影：${moviesDF.first()(1)}，年份：${moviesDF.first()(2)}")
    // head()等价于first()方法；head(n)返回数据集中前n条数据，以Array形式
    moviesDF.head(5)
    // take(n)返回数据集中前n条数据，以Array形式
    moviesDF.take(5)
    // takeAsList(n)返回数据集中前n条数据，以List形式
    moviesDF.takeAsList(5)
    // collect()返回一个包含数据集中所有行的数组
    moviesDF.collect()
    // collectAsList()返回一个包含数据集中所有行的数组，以List形式
    moviesDF.collectAsList()
    // columns返回数据集的列名，以Array形式
    println(moviesDF.columns.mkString(","))
    // dtypes返回数据集的数据类型，以Array形式
    moviesDF.dtypes.foreach(println)


    spark.stop()
  }
}
